"""
知识库服务
"""
import os
from typing import List, Dict, Any, Optional
from pathlib import Path
from langchain.docstore.document import Document

from ..config.settings import settings
from ..core.document_loader import DocumentLoader, KnowledgeFile
from ..core.text_splitter import TextSplitter
from ..core.vector_store import VectorStoreFactory


class KnowledgeBaseService:
    """知识库服务"""
    
    def __init__(
        self, 
        knowledge_base_name: str,
        vector_store_type: str = None,
        embedding_model: str = None,
        **kwargs
    ):
        self.kb_name = knowledge_base_name
        self.vector_store_type = vector_store_type or settings.vector_store.type
        self.embedding_model = embedding_model or settings.embedding.model_name
        
        # 初始化组件
        self.document_loader = DocumentLoader()
        self.text_splitter = TextSplitter()
        
        # 创建向量存储
        self.vector_store = VectorStoreFactory.create_vector_store(
            store_type=self.vector_store_type,
            embedding_model=self.embedding_model,
            **kwargs
        )
        
        # 设置路径
        self.kb_path = self._get_kb_path()
        self.doc_path = self._get_doc_path()
        self.vs_path = self._get_vs_path()
        
        # 确保目录存在
        os.makedirs(self.kb_path, exist_ok=True)
        os.makedirs(self.doc_path, exist_ok=True)
        os.makedirs(self.vs_path, exist_ok=True)
    
    def _get_kb_path(self) -> str:
        """获取知识库路径"""
        return os.path.join(settings.vector_store.storage_path, self.kb_name)
    
    def _get_doc_path(self) -> str:
        """获取文档路径"""
        return os.path.join(self.kb_path, "content")
    
    def _get_vs_path(self) -> str:
        """获取向量存储路径"""
        return os.path.join(self.kb_path, "vector_store")
    
    def add_document(
        self, 
        file_path: str, 
        chunk_size: int = None,
        chunk_overlap: int = None,
        splitter_name: str = None,
        **kwargs
    ) -> Dict[str, Any]:
        """添加文档到知识库"""
        try:
            # 创建知识库文件对象
            kb_file = KnowledgeFile(file_path, self.kb_name)
            
            # 加载文档
            docs = kb_file.file2docs()
            if not docs:
                return {"success": False, "error": "文档加载失败"}
            
            # 文本分割
            chunk_size = chunk_size or settings.text_splitter.chunk_size
            chunk_overlap = chunk_overlap or settings.text_splitter.chunk_overlap
            splitter_name = splitter_name or settings.text_splitter.splitter_name
            
            split_docs = self.text_splitter.split_documents(
                docs,
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
                splitter_name=splitter_name,
            )
            
            if not split_docs:
                return {"success": False, "error": "文档分割失败"}
            
            # 添加元数据
            for doc in split_docs:
                doc.metadata.setdefault("source", kb_file.file_path)
                doc.metadata.setdefault("kb_name", self.kb_name)
            
            # 添加到向量存储
            doc_ids = self.vector_store.add_documents(split_docs)
            
            return {
                "success": True,
                "doc_count": len(split_docs),
                "doc_ids": doc_ids,
                "file_info": kb_file.get_file_info()
            }
            
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def add_documents(
        self, 
        file_paths: List[str], 
        **kwargs
    ) -> List[Dict[str, Any]]:
        """批量添加文档"""
        results = []
        for file_path in file_paths:
            result = self.add_document(file_path, **kwargs)
            results.append(result)
        return results
    
    def search(
        self, 
        query: str, 
        top_k: int = None,
        score_threshold: float = None,
        **kwargs
    ) -> List[Dict[str, Any]]:
        """搜索相似文档"""
        top_k = top_k or settings.search.top_k
        score_threshold = score_threshold or settings.search.score_threshold
        
        try:
            # 执行搜索
            docs_and_scores = self.vector_store.similarity_search(
                query=query,
                k=top_k,
                score_threshold=score_threshold
            )
            
            # 格式化结果
            results = []
            for doc, score in docs_and_scores:
                results.append({
                    "content": doc.page_content,
                    "metadata": doc.metadata,
                    "score": score
                })
            
            return results
            
        except Exception as e:
            print(f"搜索失败: {e}")
            return []
    
    def delete_document(self, file_path: str) -> bool:
        """删除文档"""
        try:
            # 这里需要根据文件路径找到对应的文档ID
            # 简化实现：清空向量存储
            return self.vector_store.clear()
        except Exception as e:
            print(f"删除文档失败: {e}")
            return False
    
    def clear(self) -> bool:
        """清空知识库"""
        try:
            # 清空向量存储
            self.vector_store.clear()
            
            # 删除文档目录
            if os.path.exists(self.doc_path):
                import shutil
                shutil.rmtree(self.doc_path)
                os.makedirs(self.doc_path)
            
            return True
        except Exception as e:
            print(f"清空知识库失败: {e}")
            return False
    
    def get_info(self) -> Dict[str, Any]:
        """获取知识库信息"""
        return {
            "kb_name": self.kb_name,
            "vector_store_type": self.vector_store_type,
            "embedding_model": self.embedding_model,
            "kb_path": self.kb_path,
            "doc_path": self.doc_path,
            "vs_path": self.vs_path,
        }
    
    def list_documents(self) -> List[Dict[str, Any]]:
        """列出所有文档"""
        try:
            # 这里应该从向量存储中获取文档列表
            # 简化实现：返回空列表
            return []
        except Exception as e:
            print(f"获取文档列表失败: {e}")
            return [] 